# wenyan-plus-plus-pp/algorithms/vqe.py

import numpy as np
from qiskit.circuit.library import TwoLocal
from qiskit.primitives import Estimator
from qiskit.quantum_info import SparsePauliOp

from ..core.hybrid_engine import HybridEngine

class VQESolver:
    """变分量子本征求解器(VQE)实现"""
    
    def __init__(self, Hamiltonian, num_qubits, reps=1):
        """
        初始化VQE求解器
        
        参数:
            Hamiltonian: 要求解的哈密顿量
            num_qubits: 量子比特数
            reps: 电路重复次数
        """
        self.Hamiltonian = Hamiltonian
        self.num_qubits = num_qubits
        self.reps = reps
        
        # 创建评估器
        self.estimator = Estimator()
        
        # 创建混合引擎
        self.engine = HybridEngine(
            self.quantum_evaluator,
            self.create_ansatz
        )
    
    def quantum_evaluator(self, params):
        """量子评估函数"""
        # 创建ansatz电路
        qc = self.create_ansatz(params)
        
        # 计算期望值
        job = self.estimator.run(qc, self.Hamiltonian)
        result = job.result()
        
        return result.values[0]
    
    def create_ansatz(self, params):
        """创建ansatz电路"""
        # 使用TwoLocal电路
        circuit = TwoLocal(
            self.num_qubits, 
            'ry', 
            'cz', 
            reps=self.reps, 
            entanglement='linear',
            insert_barriers=True
        )
        
        if params is not None:
            # 检查参数数量是否匹配
            if len(params) != circuit.num_parameters:
                raise ValueError(
                    f"参数数量不匹配: 电路需要 {circuit.num_parameters} 个参数, "
                    f"但传入了 {len(params)} 个参数"
                )
            
            # 使用字典格式绑定参数
            param_dict = {param: value for param, value in zip(circuit.parameters, params)}
            circuit = circuit.assign_parameters(param_dict)
        
        return circuit

def create_ansatz(params, num_qubits, reps=1):
    """创建ansatz电路的辅助函数"""
    solver = VQESolver(None, num_qubits, reps)
    return solver.create_ansatz(params)